import time
from Env.client import EnvClient
from rule_player.blue_player import BluePlayer
from rule_player.red_player import RedPlayer
from generate import Generator
#from render2d.draw import draw

class Environment:
    def __init__(self, config):
        print('---------- Environment is Launching ----------')
        scen_name = config.get('scen_name', '/home/6.7(2).ntedt')
        image_name = config.get('image_name', 'registry.inspir.ai:5000/lzy-combat:1.1')
        env_id = config.get('server_host', -1)
        base_port = config.get('base_port', 6100)
        self.client =  EnvClient(env_id, base_port, image_name, scen_name)
        self.spd = config.get('step_per_decision', 1)
        self.tps = config.get('time_per_step', 10)
        print('---------- Environment Launch Finished ----------')
        self._construct_rule_player()
        self.steps = 1
        self.render_tool = None
        #self.render_tool = draw([0,6000],[0, 4000], 1500)
        
    def _construct_rule_player(self):
        blue_player = BluePlayer()
        red_player = RedPlayer()
        self.env_players = [blue_player, red_player]

    def step(self, commands):
        t0 = time.time()
        self.steps+=1
        if len(commands):
            self.client.take_action(commands)
        for i in range(self.spd):
            for player in self.env_players:
                player.set_obs(self.obs)
                cmds = player.step()
                self.client.take_action(cmds)
            target_time = self.base_time + self.tps*(self.steps*self.spd+i)
            ctime = self.client.run_step(target_time)
        t1 = time.time()
        print(f'now steps: {self.steps}, target_time:{target_time//3600}:{target_time%3600//60}:{target_time%60}  sim_time: {ctime//3600}:{ctime%3600//60}:{ctime%60} spend time : {int(t1-t0)}')

    def reset(self):
        self.steps = 0
        while not self.client.reset():
            print('Reset Failed')
            pass
        #self.client.speedup(0)
        self._construct_rule_player()
        base_time = self.client.run_step(10)
        self.base_time = (base_time // 10+1) * 10 


    def receive_obs(self):
        obs = self.client.get_observation()
        self.obs = obs
        return obs

    def is_done(self):
        return self.steps >= 120
        return False

###################################################################################

def run_episode(env, gen):
    env.reset()
    while True:
        obs = env.receive_obs()
        commands = gen.generate(obs)
        env.step(commands)
        if env.is_done():
            break



if "__main__" == __name__ :
    config = {
        'max_step': 80,
        'image_name': 'registry.inspir.ai:5000/lzy-combat:2.2',
        'abs': 'test'
    }
    env = Environment(config)
    gen = Generator()
    while True:
        run_episode(env, gen)
